import numpy as np

# -------------------------------------------------
# 工具：向上取整到最近的 2 的幂
# -------------------------------------------------
def next_pow2(n):
    return 1 << (n - 1).bit_length()

# -------------------------------------------------
# 工具：把矩阵补 0 到 n×n
# -------------------------------------------------
def pad_to_square(a, n):
    m, k = a.shape
    if m == n and k == n:
        return a
    out = np.zeros((n, n), dtype=a.dtype)
    out[:m, :k] = a
    return out

# -------------------------------------------------
# 工具：裁剪回真实尺寸
# -------------------------------------------------
def crop(a, M, N):
    return a[:M, :N]

# -------------------------------------------------
# Strassen 递归核心（方阵，n 是 2 的幂）
# -------------------------------------------------
def strassen_square(x, y, leaf=64):
    n = x.shape[0]
    if n <= leaf:          # 小到一定规模就朴素乘法
        return x @ y

    mid = n // 2
    # 分块
    A, B, C, D = x[:mid, :mid], x[:mid, mid:], x[mid:, :mid], x[mid:, mid:]
    E, F, G, H = y[:mid, :mid], y[:mid, mid:], y[mid:, :mid], y[mid:, mid:]

    # 7 个乘积
    P1 = strassen_square(A, F - H, leaf)
    P2 = strassen_square(A + B, H, leaf)
    P3 = strassen_square(C + D, E, leaf)
    P4 = strassen_square(D, G - E, leaf)
    P5 = strassen_square(A + D, E + H, leaf)
    P6 = strassen_square(B - D, G + H, leaf)
    P7 = strassen_square(A - C, E + F, leaf)

    # 4 个结果子块
    UL = P5 + P4 - P2 + P6
    UR = P1 + P2
    LL = P3 + P4
    LR = P1 + P5 - P3 - P7

    # 合并
    top = np.hstack([UL, UR])
    bot = np.hstack([LL, LR])
    return np.vstack([top, bot])

# -------------------------------------------------
# 任意形状接口：A[M,K] · B[K,N]
# -------------------------------------------------
def strassen_matmul(A: np.ndarray, B: np.ndarray, leaf=64):
    assert A.shape[1] == B.shape[0]
    M, K = A.shape
    K2, N = B.shape
    assert K == K2

    # 目标方阵尺寸
    dim = next_pow2(max(M, K, N))
    # 补 0
    AA = pad_to_square(A, dim)
    BB = pad_to_square(B, dim)
    # Strassen
    CC = strassen_square(AA, BB, leaf)
    # 裁剪
    return crop(CC, M, N)

# -------------------------------------------------
# 简单测试
# -------------------------------------------------
if __name__ == "__main__":
    np.random.seed(0)
    M, K, N = 123, 456, 789
    A = np.random.rand(M, K).astype(np.float32)
    B = np.random.rand(K, N).astype(np.float32)

    C1 = A @ B                 # 库实现
    C2 = strassen_matmul(A, B) # 自写 Strassen

    print("最大绝对误差:", np.max(np.abs(C1 - C2)))
    # 通常 <1e-5（float32 累加误差量级）